import backtrader as bt
from datetime import datetime
from strategies.utils import BaoStockPandasData
from strategies.utils.Datas import get_stock_data2

class ChannelBreakoutStrategy(bt.Strategy):
    params = (
        ('period', 20),  # 计算通道突破的周期长度
        ('pct_change', 0.01),  # 价格变动百分比作为突破阈值
        ('stop_loss', 0.05),  # 止损比例
        ('take_profit', 0.10),  # 止盈比例
    )

    def __init__(self):
        self.data_close = self.datas[0].close
        self.high_breakout = None
        self.low_breakout = None
        self.order = None
        self.buy_price = None
        self.buy_comm = None

    def next(self):
        if self.order:
            return

        # 检查数据长度是否足够
        #if len(self.data_close) < self.params.period:
        #    return  # 数据不足，跳过当前周期

        # 计算周期内的最高价和最低价
        period_high = max(self.data_close.get(size=self.params.period))
        period_low = min(self.data_close.get(size=self.params.period))
        
        # 打印调试信息
        print(f"Current Close: {self.data_close[0]}, Period High: {period_high}, Period Low: {period_low}")

        # 检查是否突破前期高点
        if not self.position and self.data_close[0] > period_high * (1 + self.params.pct_change):
            self.log(f"Long Breakout at {self.data_close[0]}")
            self.high_breakout = self.data_close[0]
            amount_to_invest = (0.95 * self.broker.cash) / self.data_close[0]
            self.buy(size=amount_to_invest)
            self.buy_price = self.data_close[0]
            self.buy_comm = self.broker.getcommission()

        # 检查是否突破前期低点
        elif not self.position and self.data_close[0] < period_low * (1 - self.params.pct_change):
            self.log(f"Short Breakout at {self.data_close[0]}")
            self.low_breakout = self.data_close[0]
            amount_to_invest = (0.95 * self.broker.cash) / self.data_close[0]
            self.sell(size=amount_to_invest)
            self.buy_price = self.data_close[0]
            self.buy_comm = self.broker.getcommission()

        # 止损
        if self.position and self.data_close[0] <= self.buy_price * (1 - self.params.stop_loss):
            self.log(f"Stop Loss at {self.data_close[0]}")
            self.close()

        # 止盈
        elif self.position and self.data_close[0] >= self.buy_price * (1 + self.params.take_profit):
            self.log(f"Take Profit at {self.data_close[0]}")
            self.close()

    def log(self, txt, dt=None):
        ''' Logging function for this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return  # 订单状态未最终确定，等待

        if order.status in [order.Completed]:
            if order.isbuy():
                self.log(f"BUY EXECUTED --- Price: {order.executed.price}, Cost: {order.executed.value}, Commission: {order.executed.comm}")
            else:  # Sell
                self.log(f"SELL EXECUTED --- Price: {order.executed.price}, Size: {order.executed.size}, Commission: {order.executed.comm}")

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

# 回测设置示例
if __name__ == '__main__':
    cerebro = bt.Cerebro()

    # 假设已经通过某种方式（如Yahoo Finance）获取了数据
    df = get_stock_data2("sh.600519", "2005-01-01", "2023-12-31")
    data = BaoStockPandasData()
    cerebro.adddata(data)

    cerebro.addstrategy(ChannelBreakoutStrategy, period=5, pct_change=0.01)

    cerebro.broker.setcash(100000.0)
    cerebro.broker.set_coc(True)  # 现金交易模式

    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())